import sys
import os
import subprocess
import time
import threading
import configs

nougat_conf = configs.Nougat()
base_conf = configs.Base()
llm_conf = configs.LLM()

if not os.path.exists(base_conf.log_dir):
    os.makedirs(base_conf.log_dir)

if not os.path.exists(base_conf.data_dir):
    os.makedirs(base_conf.data_dir)

def run_nougat():
    env = os.environ.copy()
    extra_env = dict(CUDA_VISIBLE_DEVICES=nougat_conf.gpu_id,
                     NOUGAT_CHECKPOINT=nougat_conf.model_path,
                     NOUGAT_BATCHSIZE=str(nougat_conf.batch_size))
    env.update(extra_env)
    with open(f'{base_conf.log_dir}/nougat.log', 'w') as fout:
        extra_env_str = ' '.join(f'{k}={v}' for k, v in extra_env.items())
        cmd = 'python3 nougat_api.py --host 0.0.0.0 --port 8503'
        print(f'[CMD] {extra_env_str} {cmd}', file=sys.stderr)
        subprocess.run(cmd.split(), env=env, stdout=fout, stderr=subprocess.STDOUT)

def run_controller():
    with open(f'{base_conf.log_dir}/controller.log', 'w') as fout:
        cmd = 'python3 -m fastchat.serve.controller --host 127.0.0.1 --port 8760'
        print(f'[CMD] {cmd}', file=sys.stderr)
        subprocess.run(cmd.split(), stdout=fout, stderr=subprocess.STDOUT)

def run_model_worker():
    env = os.environ.copy()
    extra_env = dict(CUDA_VISIBLE_DEVICES=llm_conf.gpu_id)
    env.update(extra_env)
    with open(f'{base_conf.log_dir}/model_worker.log', 'w') as fout:
        extra_env_str = ' '.join(f'{k}={v}' for k, v in extra_env.items())
        cmd = f'python3 -m fastchat.serve.model_worker --controller-address http://127.0.0.1:8760 --model-path {llm_conf.model_path} --host 127.0.0.1 --port 8761 --worker-address http://127.0.0.1:8761'
        print(f'[CMD] {extra_env_str} {cmd}')
        subprocess.run(cmd.split(), env=env, stdout=fout, stderr=subprocess.STDOUT)

def run_openai_api():
    with open(f'{base_conf.log_dir}/openai_api.log', 'w') as fout:
        cmd = 'python3 -m fastchat.serve.openai_api_server --controller-address http://127.0.0.1:8760 --host 0.0.0.0 --port 8765'
        print(f'[CMD] {cmd}', file=sys.stderr)
        subprocess.run(cmd.split(), stdout=fout, stderr=subprocess.STDOUT)

nougat_thread = threading.Thread(target=run_nougat)
nougat_thread.start()

os.environ['PYTHONPATH'] = 'fastchat:' + os.environ['PYTHONPATH']
os.environ['LOGDIR'] = ''

controller_thread = threading.Thread(target=run_controller)
controller_thread.start()

time.sleep(1)

model_worker_thread = threading.Thread(target=run_model_worker)
model_worker_thread.start()

openai_api_thread = threading.Thread(target=run_openai_api)
openai_api_thread.start()

nougat_thread.join()
controller_thread.join()
model_worker_thread.join()
openai_api_thread.join()
